# -*- coding: utf-8 -*-
import time
import threadpool
import tushare as ts
import pymongo
from apscheduler.schedulers.blocking import BlockingScheduler
# from datetime import datetime, date, timedelta
from mongo_config import *
mongo_client = pymongo.MongoClient(MONGO_URL_ALI, 27017)
db = mongo_client[MONGO_DB_RES]

def get_5minute(code):
    re_data = {}
    dictlist = []
    stock = code + '_5'
    data = ts.get_k_data(code, ktype='5', pause=1)[-3:-1]
    for index, datum in data.iterrows():
        dicts = {}
        dicts['_id'] = int(time.mktime(time.strptime(datum['date'], "%Y-%m-%d %H:%M")))
        dicts['open'] = datum['open']
        dicts['close'] = datum['close']
        dicts['high'] = datum['high']
        dicts['low'] = datum['low']
        dicts['volume'] = datum['volume']
        dictlist.append(dicts)
    re_data[stock] = dictlist
    if len(re_data) > 0:
        save_to_mongo(re_data)


def get_15minute(code):
    re_data = {}
    dictlist = []
    stock = code + '_15'
    data = ts.get_k_data(code, ktype='15', pause=1)[-3:-1]
    for index, datum in data.iterrows():
        dicts = {}
        dicts['_id'] = int(time.mktime(time.strptime(datum['date'], "%Y-%m-%d %H:%M")))
        dicts['open'] = datum['open']
        dicts['close'] = datum['close']
        dicts['high'] = datum['high']
        dicts['low'] = datum['low']
        dicts['volume'] = datum['volume']
        dictlist.append(dicts)
    re_data[stock] = dictlist
    if len(re_data) > 0:
        save_to_mongo(re_data)


def get_30minute(code):
    re_data = {}
    dictlist = []
    stock = code + '_30'
    data = ts.get_k_data(code, ktype='30', pause=1)[-2:-1]
    for index, datum in data.iterrows():
        dicts = {}
        dicts['_id'] = int(time.mktime(time.strptime(datum['date'], "%Y-%m-%d %H:%M")))
        dicts['open'] = datum['open']
        dicts['close'] = datum['close']
        dicts['high'] = datum['high']
        dicts['low'] = datum['low']
        dicts['volume'] = datum['volume']
        dictlist.append(dicts)
    re_data[stock] = dictlist
    if len(re_data) > 0:
        save_to_mongo(re_data)
        time.sleep(1)


def get_1h(code):
    re_data = {}
    dictlist = []
    stock = code + '_1h'
    data = ts.get_k_data(code, ktype='60', pause=1)[-2:-1]
    for index, datum in data.iterrows():
        dicts = {}
        dicts['_id'] = int(time.mktime(time.strptime(datum['date'], "%Y-%m-%d %H:%M")))
        dicts['open'] = datum['open']
        dicts['close'] = datum['close']
        dicts['high'] = datum['high']
        dicts['low'] = datum['low']
        dicts['volume'] = datum['volume']
        dictlist.append(dicts)
    re_data[stock] = dictlist
    if len(re_data) > 0:
        save_to_mongo(re_data)
        time.sleep(1)


def get_1D(code):
    re_data = {}
    dictlist = []
    stock = code + '_1D'
    data = ts.get_k_data(code, ktype='D', pause=3)[-1:]
    for index, datum in data.iterrows():
        dicts = {}
        dicts['_id'] = int(time.mktime(time.strptime(datum['date'], "%Y-%m-%d")))
        dicts['open'] = datum['open']
        dicts['close'] = datum['close']
        dicts['high'] = datum['high']
        dicts['low'] = datum['low']
        dicts['volume'] = datum['volume']
        dictlist.append(dicts)
    re_data[stock] = dictlist
    if len(re_data) > 0:
        save_to_mongo(re_data)
        time. sleep(3)


def get_1W(code):
    re_data = {}
    dictlist = []
    stock = code + '_1W'
    data = ts.get_k_data(code, ktype='W', pause=3)[-1:]
    for index, datum in data.iterrows():
        dicts = {}
        dicts['_id'] = int(time.mktime(time.strptime(datum['date'], "%Y-%m-%d")))
        dicts['open'] = datum['open']
        dicts['close'] = datum['close']
        dicts['high'] = datum['high']
        dicts['low'] = datum['low']
        dicts['volume'] = datum['volume']
        dictlist.append(dicts)
    re_data[stock] = dictlist
    if len(re_data) > 0:
        save_to_mongo(re_data)
        time.sleep(10)


def get_1M(code):
    re_data = {}
    dictlist = []
    stock = code + '_1M'
    data = ts.get_k_data(code, ktype='M', pause=3)[-1:]
    for index, datum in data.iterrows():
        dicts = {}
        dicts['_id'] = int(time.mktime(time.strptime(datum['date'], "%Y-%m-%d")))
        dicts['open'] = datum['open']
        dicts['close'] = datum['close']
        dicts['high'] = datum['high']
        dicts['low'] = datum['low']
        dicts['volume'] = datum['volume']
        dictlist.append(dicts)
    re_data[stock] = dictlist
    if len(re_data) > 0:
        save_to_mongo(re_data)
        time.sleep(10)

def save_to_mongo(datas):
    # print(datas)
    key = [key for key in datas][0]
    values = datas[key]
    collection_res = db[key]
    try:
        collection_res.insert_many(values, ordered=False)
        time.sleep(0.1)
        print(key, 'insert ok')
    except Exception as e:
        print(e)

def main_5():
    stkcodes = stk_codes
    start_time = time.time()
    pool = threadpool.ThreadPool(15)
    for i in range(0,len(stkcodes),100):
        codes = stkcodes[i:i+100]
        requests = threadpool.makeRequests(get_5minute, codes)
        [pool.putRequest(req) for req in requests]
        time.sleep(3)
    pool.wait()
    print('%d second' % (time.time() - start_time))


def main_15():
    stkcodes = stk_codes
    start_time = time.time()
    pool = threadpool.ThreadPool(10)
    for i in range(0, len(stkcodes), 100):
        codes = stkcodes[i:i + 100]
        requests = threadpool.makeRequests(get_15minute, codes)
        [pool.putRequest(req) for req in requests]
        time.sleep(2)
    pool.wait()
    print('%d second' % (time.time() - start_time))


def main_30():
    stkcodes = stk_codes
    start_time = time.time()
    pool = threadpool.ThreadPool(10)
    for i in range(0, len(stkcodes), 100):
        codes = stkcodes[i:i + 100]
        requests = threadpool.makeRequests(get_30minute, codes)
        [pool.putRequest(req) for req in requests]
        time.sleep(2)
    pool.wait()
    print('%d second' % (time.time() - start_time))


def main_1h():
    stkcodes = stk_codes
    start_time = time.time()
    pool = threadpool.ThreadPool(10)
    for i in range(0, len(stkcodes), 100):
        codes = stkcodes[i:i + 100]
        requests = threadpool.makeRequests(get_1h, codes)
        [pool.putRequest(req) for req in requests]
        time.sleep(5)
    pool.wait()
    print('%d second' % (time.time() - start_time))


def main_1D():
    stkcodes = stk_codes
    start_time = time.time()
    pool = threadpool.ThreadPool(12)
    for i in range(0, len(stkcodes), 100):
        codes = stkcodes[i:i + 100]
        requests = threadpool.makeRequests(get_1D, codes)
        [pool.putRequest(req) for req in requests]
        time.sleep(3)
    pool.wait()
    print('%d second' % (time.time() - start_time))


def main_1W():
    stkcodes = stk_codes
    start_time = time.time()
    pool = threadpool.ThreadPool(10)
    for i in range(0, len(stkcodes), 100):
        codes = stkcodes[i:i + 100]
        requests = threadpool.makeRequests(get_1W, codes)
        [pool.putRequest(req) for req in requests]
        time.sleep(10)
    pool.wait()
    print('%d second' % (time.time() - start_time))


def main_1M():
    stkcodes = stk_codes
    start_time = time.time()
    pool = threadpool.ThreadPool(10)
    for i in range(0, len(stkcodes), 100):
        codes = stkcodes[i:i + 100]
        requests = threadpool.makeRequests(get_1M, codes)
        [pool.putRequest(req) for req in requests]
        time.sleep(10)
    pool.wait()
    print('%d second' % (time.time() - start_time))

if __name__ == '__main__':
    # main_1M()
    # main_1M()
    # 定义BlockingScheduler
    # '''
    sched = BlockingScheduler()
    print('start:', time.localtime())
    sched.add_job(main_5, 'cron', day_of_week='mon-fri', hour='9-11,13-14',minute='1,6,11,16,21,26,31,36,41,46,51,56')
    sched.add_job(main_15, 'cron', day_of_week='mon-fri', hour='9-11,13-14',minute='8,23,38,42,57')
    sched.add_job(main_30, 'cron', day_of_week='mon-fri', hour='9-11,13-14',minute='3,33')
    sched.add_job(main_1h, 'cron', day_of_week='mon-fri', hour='9-11,13-14',minute='13')
    sched.add_job(main_1D, 'cron', day_of_week='mon-fri', hour=15, minute=5)
    sched.add_job(main_1W, 'cron', day_of_week='fri', hour=18, minute=30)
    sched.add_job(main_1M, 'cron', month='1-12', day=2, hour=2, minute=30)
    sched.start()
    # '''
    # 休市后追加最后的数据
    # sched.add_job(main_30, 'cron', day_of_week='mon-fri',hour=15, minute=5)
    # sched.add_job(main_1h, 'cron', day_of_week='mon-fri',hour=15, minute=20)
    # sched.add_job(main_1D, 'cron', hour=22, minute=0)
    # sched.add_job(main_1W, 'cron', hour=23, minute=30)
    # sched.add_job(main_1M, 'cron', hour=0, minute=30)
